# -*- coding:utf-8 -*-
'''
@paper:Different applied median filter in salt and pepper noise
'''
import numpy as np

# Different applied median filter
class Damf(object):
    def binary_matrix(self, A):
        '''Set the noise entry to 0.Set regular entry to 1'''
        B = np.zeros(np.shape(A))
        B[(A == 255) | (A == 0)] = 0
        B[(A != 255) & (A != 0)] = 1
        return B

    def t_symmetric_pad_matrix(self, A, t):
        m, n = np.shape(A)
        img_pad = np.zeros((m + t, n + t))
        img_pad[t:, t:] = A
        left_up = np.rot90(np.rot90(A[:t, :t]))
        up = np.flip(A[:t, :], 0)
        right_up = np.rot90(np.rot90(A[:t, -t:]))

        left = np.flip(A[:, 0:t], 1)
        right = np.flip(A[:, -t:], 1)

        left_down = np.rot90(np.rot90(A[-t:, :t]))
        dowm = np.flip(A[-t:, :], 0)
        right_down = np.rot90(np.rot90(A[-t:, -t:]))

        up_h = np.hstack((left_up, up, right_up))  # h：horizontal
        middle_h = np.hstack((left, A, right))
        down_h = np.hstack((left_down, dowm, right_down))

        img_pad = np.vstack((up_h, middle_h, down_h))
        return img_pad

    def k_approximate_matrix(self, img_pad, i, j, k, t):
        k_approximate_mat = img_pad[(i + t - k):(i + t + k + 1), (j + t - k):(j + t + k + 1)]
        return k_approximate_mat

    def regular_row_matrix(self, img_k_approximate):
        H = img_k_approximate.flatten()
        H = np.sort(H)
        R = H[(H != 255) & (H != 0)]

        return R

    def process_image(self, nim):
        t1 = 3
        img_DAMF = nim.copy()
        B = self.binary_matrix(nim)
        P_A_3 = self.t_symmetric_pad_matrix(nim, t1)
        P_B_3 = self.t_symmetric_pad_matrix(B, t1)
        m, n = np.shape(nim)
        for i in range(m):
            for j in range(n):
                if B[i, j] == 0:
                    B_ij_1 = self.k_approximate_matrix(P_B_3, i, j, 1, t1)
                    B_ij_2 = self.k_approximate_matrix(P_B_3, i, j, 2, t1)
                    B_ij_3 = self.k_approximate_matrix(P_B_3, i, j, 3, t1)
                    if B_ij_1.any() > 0:  # 使用3x3板
                        A_ij_1 = self.k_approximate_matrix(P_A_3, i, j, 1, t1)
                        R_ij_1 = self.regular_row_matrix(A_ij_1)
                        mR_ij_1 = np.median(R_ij_1)
                        img_DAMF[i, j] = mR_ij_1

                    elif B_ij_2.any() > 0:  # 使用5x5模板
                        A_ij_2 = self.k_approximate_matrix(P_A_3, i, j, 2, t1)
                        R_ij_2 = self.regular_row_matrix(A_ij_2)
                        mR_ij_2 = np.median(R_ij_2)
                        img_DAMF[i, j] = mR_ij_2

                    elif B_ij_3.any() > 0:  # 使用7x7模板
                        A_ij_3 = self.k_approximate_matrix(P_A_3, i, j, 3, t1)
                        R_ij_3 = self.regular_row_matrix(A_ij_3)
                        mR_ij_3 = np.median(R_ij_3)
                        img_DAMF[i, j] = mR_ij_3
        B = self.binary_matrix(img_DAMF)
        t2 = 1
        P_A_1 = self.t_symmetric_pad_matrix(img_DAMF, t2)
        P_B_1 = self.t_symmetric_pad_matrix(B, t2)
        for i in range(m):
            for j in range(n):
                if B[i, j] == 0:
                    B_ij_1 = self.k_approximate_matrix(P_B_1, i, j, 1, t2)
                    if B_ij_1.any() > 0:
                        A_ij_1 = self.k_approximate_matrix(P_A_1, i, j, 1, t2)
                        R_ij_1 = self.regular_row_matrix(A_ij_1)
                        mR_ij_1 = np.median(R_ij_1)
                        img_DAMF[i, j] = mR_ij_1

        return img_DAMF
